from typing import Tuple, Literal, List, Optional, TypedDict

import torch
import numpy as np
from torch.utils.data import Dataset

from preprocessing import text_preprocess, EncodingData

class TelecomFraudData(TypedDict):
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    labels: torch.Tensor

class TelecomFraudDataset(Dataset):
    _preprocessing_data: Optional[Tuple[List[EncodingData], List[EncodingData]]] = None

    def __init__(self, mode: Literal['train', 'test']):
        super().__init__()

        if TelecomFraudDataset._preprocessing_data is None:
            TelecomFraudDataset._preprocessing_data = text_preprocess()

        assert mode in ['train', 'test']
        self.data = TelecomFraudDataset._preprocessing_data[0 if mode == 'train' else 1]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index: int) -> TelecomFraudData:
        encoding_data = self.data[index]

        return {
            'input_ids': torch.tensor(encoding_data.input_ids, dtype=torch.int32),
            'attention_mask': torch.tensor(encoding_data.attention_mask, dtype=torch.int32),
            'labels': torch.tensor(np.eye(5)[encoding_data.label], dtype=torch.float32)
        }


if __name__ == '__main__':
    dataset = TelecomFraudDataset('train')
    print(dataset[0])